import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split
import lightgbm as lgb
from sqlalchemy import create_engine


def fetch_data_from_db():
    """从数据库获取BWIC和投标数据"""
    db_url = "mysql+pymysql://root:yourpassword@localhost:3306/salessvc?charset=utf8mb4"
    engine = create_engine(db_url, pool_pre_ping=True)

    try:
        # 获取BWIC摘要数据
        df_summary = pd.read_sql("SELECT * FROM auction_assets", engine)

        # 获取投标详情数据
        df_detail = pd.read_sql("SELECT * FROM auction_bids", engine)

        # 合并数据
        df_res = df_summary.merge(df_detail, how='inner', left_on='id', right_on='auction_id')
        df_res.drop(['id', 'auction_id'], axis=1, inplace=True)

        print(f"获取到BWIC记录数: {df_summary.shape[0]}")
        print(f"获取到投标记录数: {df_detail.shape[0]}")
        print(f"合并后记录数: {df_res.shape[0]}")

        return df_res

    except Exception as e:
        print(f"数据库查询错误: {e}")
        return pd.DataFrame()
    finally:
        engine.dispose()


def preprocess_data(df: pd.DataFrame):
    """数据预处理"""
    # 重置索引
    df.reset_index(inplace=True)

    # 分类变量转换
    cat_columns = ['bond_code', 'sett_country', 'issuer', 'bidder_id', 'client_id', 'bid_successful']
    for col in cat_columns:
        if col in df.columns:
            df[col] = df[col].astype('category')

    # 处理时间字段
    if 'created_time' in df.columns:
        df['created_time'] = pd.to_datetime(df['created_time'])
        df['hour'] = df['created_time'].dt.hour
        df['day_of_week'] = df['created_time'].dt.dayofweek

    return df


def get_lgb_params():
    """LightGBM模型参数"""
    return {
        "objective": "regression",
        "metric": "rmse",
        "num_leaves": 60,
        "max_depth": -1,
        "learning_rate": 0.01,
        "bagging_fraction": 0.9,
        "feature_fraction": 0.9,
        "bagging_freq": 5,
        "bagging_seed": 2018,
        "verbosity": -1,
        "early_stopping_rounds": 100
    }


def train_and_evaluate(df: pd.DataFrame, target_col: str = "starting_amount"):
    """训练和评估模型"""
    # 划分训练集和测试集
    train, test = train_test_split(df, test_size=0.2, random_state=42)

    # 准备数据
    features = [col for col in df.columns if col not in [target_col, 'index', 'created_time']]
    train_x = train[features]
    test_x = test[features]
    train_y = train[target_col]
    test_y = test[target_col]

    # 创建数据集
    train_ds = lgb.Dataset(train_x, label=train_y)
    valid_ds = lgb.Dataset(test_x, label=test_y, reference=train_ds)

    # 训练模型
    params = get_lgb_params()
    model = lgb.train(
        params,
        train_ds,
        num_boost_round=1000,
        valid_sets=[train_ds, valid_ds],
        verbose_eval=50
    )

    # 评估模型
    pred_y = model.predict(test_x)
    rmse = np.sqrt(mean_squared_error(test_y, pred_y))
    mae = mean_absolute_error(test_y, pred_y))
    r2 = r2_score(test_y, pred_y))

    print(f"\n模型评估结果:")
    print(f"RMSE: {rmse:.4f}")
    print(f"MAE: {mae:.4f}")
    print(f"R2 Score: {r2:.4f}")

    # 可视化结果
    plt.figure(figsize=(15, 6))
    plt.plot(test_y.values[:100], 'r-', label='Actual')
    plt.plot(pred_y[:100], 'g--', label='Predicted')
    plt.title('Actual vs Predicted Values (First 100 Samples)')
    plt.legend()
    plt.show()

    # 特征重要性
    feat_imp = pd.Series(
        model.feature_importance(),
        index=features
    ).sort_values(ascending=False)

    plt.figure(figsize=(10, 8))
    feat_imp.head(20).plot(kind='barh')
    plt.title('Top 20 Important Features')
    plt.show()

    return model, feat_imp


def save_predictions(model, test_data, original_test):
    """保存预测结果"""
    predictions = model.predict(test_data)
    output = pd.DataFrame({
        'bond_code': original_test['bond_code'],
        'actual_amount': original_test['starting_amount'],
        'predicted_amount': predictions
    })

    # 保存到CSV
    output.to_csv('bond_price_predictions.csv', index=False)
    print("\n预测结果已保存到 bond_price_predictions.csv")

    return output


if __name__ == '__main__':
    # 1. 从数据库获取数据
    print("正在从数据库获取数据...")
    df_data = fetch_data_from_db()

    if df_data.empty:
        print("无法获取数据，程序退出")
        exit()

    # 2. 数据预处理
    print("\n数据预处理中...")
    df_processed = preprocess_data(df_data)
    print(df_processed.info())

    # 3. 训练和评估模型
    print("\n开始训练模型...")
    target_column = "starting_amount"  # 假设我们要预测起始价格
    model, feature_importance = train_and_evaluate(df_processed, target_column)

    # 4. 保存模型和预测结果
    model.save_model('bond_price_model.txt')
    print("\n模型已保存为 bond_price_model.txt")

    # 重新加载模型测试
    loaded_model = lgb.Booster(model_file='bond_price_model.txt')

    # 准备测试数据
    _, test = train_test_split(df_processed, test_size=0.2, random_state=42)
    features = [col for col in df_processed.columns if col not in [target_column, 'index', 'created_time']]
    test_x = test[features]

    # 生成预测结果
    predictions = save_predictions(loaded_model, test_x, test)
    print("\n预测结果示例:")
    print(predictions.head())